#include <sys/wait.h>
#include <sys/stat.h>
#include <stdio.h>
#include <fcntl.h>
#include <sys/types.h>
#include <string.h>
#include <alloca.h>
#include <errno.h>
#include <stdlib.h>

#include "cfd-ptrace.h"

int load_text(pid_t    pid
              , long   offset
              , void * data
              , long   data_size
              )
{
    char * current_address = data;
    int i, j;

    union __data_u
    {
        long as_long;
        char as_char[sizeof(long)];
    } __data;

    i = 0;
    j = data_size / sizeof(long);

    while(i < j)
    {
        __data.as_long = ptrace(PTRACE_PEEKDATA
                                , pid
                                , offset + i * sizeof(long)
                                , NULL
                                );

        memcpy(current_address
               , __data.as_char
               , sizeof(long)
               );

        current_address += sizeof(long)
            , i++
            ;
    }

    j = data_size % sizeof(long);

    if(j != 0)
    {
        __data.as_long = ptrace(PTRACE_PEEKDATA
                                , pid
                                , offset + i * sizeof(long)
                                , NULL
                                );
        memcpy(current_address
               , __data.as_char
               , j
               );
    }

    return 0;
}

int store_text(pid_t    pid
              , long   offset
              , void * data
              , long   data_size
              )
{
    char * current_address = data;
    int i, j;

    union __data_u
    {
        long as_long;
        char as_char[sizeof(long)];
    } __data;

    i = 0;
    j = data_size / sizeof(long);

    int rval = 0;

    while(i < j)
    {
        memcpy(__data.as_char
               , current_address
               , sizeof(long)
               );

        rval = ptrace(PTRACE_POKEDATA
                      , pid
                      , offset + i * sizeof(long)
                      , __data.as_long
                      );

        if (0 != rval)
            return rval;

        current_address += sizeof(long)
            , i++
            ;
    }

    j = data_size % sizeof(long);

    if(j != 0)
    {
        memcpy(__data.as_char
               , current_address
               , j
               );

        rval = ptrace(PTRACE_POKEDATA
                      , pid
                      , offset + i * sizeof(long)
                      , __data.as_long
                      );
    }

    return rval;
}

/* Default error handler */

void __default_error_handler (pid_t pid
                              , int signal_value
                              , struct user_regs_struct * error_registers
                              , struct user_regs_struct * orig_registers
                              , void * text
                              , long text_size
                              , long text_addr
                              )
{
    /* Just write some info */

    fprintf(stderr
            , " ============== ERROR ============= \n"
            );

    if (WSTOPSIG(signal_value != SIGSEGV))
    {
        fprintf(stderr
                , "SIGNAL : %d\n"
                  " ================================== \n"
                , WSTOPSIG(signal_value)
                );
        return ;
    }

    siginfo_t sig;

    fprintf(stderr
            , "SEGMENTATION FAULT\n"
              "CODE SIZE: %ld\n"
            , text_size
            );

    ptrace(PTRACE_GETSIGINFO, pid, NULL, &sig);

    fprintf(stderr
            , "_AX: %p\n_BX: %p\n_CX: %p\n_DX: %p\n_SP: %p ; _IP: %p\n"
            , (void *) error_registers->R_0
            , (void *) error_registers->R_1
            , (void *) error_registers->R_2
            , (void *) error_registers->R_3
            , (void *) error_registers->R_DATA
            , (void *) error_registers->R_PC
            );

    fprintf(stderr
            , "R:_IP: %ld\n"
            , error_registers->R_DATA - text_addr
            );

    fprintf(stderr
            , "Fault at address: %p\n"
              "R:SEGV ADDR: %p\n"
            , sig.si_addr
            , ((char *) sig.si_addr) - text_addr
            );

    fprintf(stderr
            ,  " ================================== \n"
            );
}


int run_text(pid_t pid
             , void          * text
             , long            text_size
             , void          * data
             , long            data_size
             , struct filter * filter
             , handle_error    handler
             )
{
    struct user_regs_struct orig_registers, new_registers;
    int signal_value;
    int rval = 1;

    long text_addr = DEFAULT_BASE;

    char * orig_text = NULL;
    char * orig_data = NULL;

    if (-1 == ptrace(PTRACE_ATTACH
                     , pid
                     , NULL
                     , NULL
                     ))
    {
        fprintf(stderr
                , "Couldn't attach: %s\n"
                , strerror(errno)
                );

        goto lbEnd2;
    }

    /* Wait for SIGCONT */
    waitpid(pid, &signal_value, 0);

    if ((WSTOPSIG(signal_value) != SIGCONT)
        && (WSTOPSIG(signal_value) != SIGCONT + 1 // 18,19
            ))

    {
        fprintf(stderr
                , "Bad signal: %d != %d\n"
                , WSTOPSIG(signal_value)
                , SIGCONT
                );

        ptrace(PTRACE_DETACH, pid, NULL, NULL);

        goto lbEnd2;
    }

    if(-1 == ptrace(PTRACE_GETREGS
                    , pid
                    , NULL
                    , &orig_registers
                    ))
    {
        fprintf(stderr
                , "Couldn't get registers: %s\n"
                , strerror(errno)
                );

        goto lbEnd2;
    }

    memcpy(&new_registers, &orig_registers, sizeof(new_registers));

    /* Make register/text/data filtration, if any */
    if(0 != apply_filter(pid
                         , & new_registers
                         , & text
                         , & text_size
                         , & text_addr
                         , & data
                         , & data_size
                         , filter
                         ))
    {
        fprintf(stderr
                , "Couldn't apply specified filter\n"
                );

        goto lbEnd2;
    }


    /* Save original text */
    orig_text  = malloc (text_size);

    if (NULL == orig_text)
    {
        fprintf(stderr
                , "Couldn't allocate memory: %s\n"
                , strerror(errno)
                );

        goto lbEnd2;
    }

    /* I'll inject code here */
    new_registers.R_PC = text_addr;

    /* Make code backup */
    load_text(pid
              , new_registers.R_PC
              , orig_text
              , text_size
              );

    /* Load new code */
    rval =
    store_text(pid
               , new_registers.R_PC
               , text
               , text_size
               );

    if (rval != 0)
    {
        fprintf(stderr
                , "Couldn't store text to process: %s\n"
                , strerror(errno)
                );

        goto lbEnd;
    }

    /* If data required - replace it with new one */
    if (data && data_size)
    {
        orig_data = malloc (data_size);

        if (orig_data == NULL)
        {
            fprintf(stderr
                    , "Couldn't allocate temporary buffer for original stack: %s\n"
                    , strerror(errno)
                    );

            goto lbEnd;
        }

        /* Make stack backup */
        load_data(pid
                  , new_registers.R_DATA
                  , orig_data
                  , data_size
                  );

        /* Place new stack */
        int rval =
            store_data(pid
                       , new_registers.R_DATA
                       , data
                       , data_size
                       );

        if (rval != 0)
        {
            fprintf(stderr
                    , "Couldn't store stack: %s\n"
                    , strerror(errno)
                    );

            goto lbEnd;
        }

        /* Increment stack with data size  */
        new_registers.R_DATA += data_size;
    }


    /* Save new R_PC/R_DATA values */
    ptrace(PTRACE_SETREGS, pid, NULL, &new_registers);

    /* Apply/Run new code */
    ptrace(PTRACE_CONT, pid, NULL, NULL);

    /* Wait for SIGTRAP. Code must call int3 at the end */
    waitpid(pid, &signal_value, 0);

    /* If signal != SIGTRAP, then some shit happens. Lets our handler handle it */
    if (WSTOPSIG(signal_value) != SIGTRAP)
    {
        fprintf(stderr
                , "Tracer catch unwaited signal: %d\n"
                , signal_value
                );

        struct user_regs_struct error_registers;

        if(-1 == ptrace(PTRACE_GETREGS
                        , pid
                        , NULL
                        , & error_registers
                        ))
        {
            fprintf(stderr
                    , "Couldn't load registers after crash. Abort: %s\n"
                    , strerror(errno));

            goto lbEnd;
        }

        if (NULL == handler)
            handler = __default_error_handler;

        handler(pid
                , signal_value
                , &error_registers
                , &orig_registers
                , text
                , text_size
                , text_addr
                );
    }
    else
        rval = 0;

    /* Anyway, we must restore orig state.
     * Hopelly, injected code do this stuff too. */

    /* Restore original register values */
    ptrace(PTRACE_SETREGS
           , pid
           , NULL
           , &orig_registers
           );

    /* Restore original code */
    store_text(pid
               , new_registers.R_PC
               , orig_text
               , text_size
               );

    /* If stack changed, restore the old one
     * at old point */
    if (orig_data)
        store_data(pid
                   , orig_registers.R_DATA
                   , orig_data
                   , data_size
                   );

  lbEnd2:
    /* Leave process */
    ptrace(PTRACE_DETACH, pid, NULL, NULL);

  lbEnd:
    free(orig_text);
    free(orig_data);
    return rval;
}


